#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/12/18 14:46
# @Author  : thinkive_cfy_ide_3
# @Site    : 
# @File    : nav_based_allocation.py
# @Software: PyCharm 

"""
脚本说明
"""
from quant_researcher.quant.datasource_fetch.factor_api import factor_return_related
from quant_researcher.quant.datasource_fetch.factor_api.factor_constant import BARRA_STYLE_FACTOR, BOND_7_FACTOR, \
    FAMA_5_FACTOR, BARRA_FULL_FACTOR
from quant_researcher.quant.datasource_fetch.index_api import index_price_related
from quant_researcher.quant.datasource_fetch.index_api.index_constant import ASSET_INDEX_DICT, SW_CODE_LIST, \
    ZZ_CODE_LIST
from quant_researcher.quant.project_tool.math_func.regression_tools import get_regression_res_by_nav


def get_nav_based_allocation(portfolio_ret, analysis_type, reg_method, need_bounds,
                             need_constraints, lasso_alpha, ridge_alpha):
    """

    :param pd.DataFrame portfolio_ret: 组合的收益时间序列
    -------------+---------------+
     trade_date |          ret        |
    -------------+---------------+
    2020-01-02   |      0.1121    |
    2020-01-03   |      0.0968    |
    2020-01-04   |      0.0192    |
    2020-01-05   |      0.0879    |
    2020-01-06   |      0.0816    |
    2020-01-07   |      0.0971    |
    2020-01-08   |      0.0162    |
    2020-01-09   |      0.0049    |
    2020-01-10   |      0.0049    |
    2020-01-11   |      0.0049    |
    2020-01-12   |      0.0049    |
    ---------------+--------------+
    :param analysis_type: asset，大类资产配置（包括股票，债券，现金）；
                                    sw28，申万行业配置（包括申万一级28行业）；
                                    zz10，中证行业配置（包括中证一级10行业）;
                                    barra，barra风格因子配置（包括barra10个风格因子）。
    :param int reg_method: 回归算法，1-最小二乘，2-lasso，3-ridge
    :param bool need_bounds: 是否需要给回归系数设置边界，最小二乘设置的边界为(0, 1)，lasso设置的边界为(0, ∞)，
                             ridge无法设置边界条件，此时也需赋值（0或1均可）
    :param bool need_constraints: 是否需要给回归系数设置约束条件，该参数只针对minimize回归，若为True，则系数之和为1，
                                  其他两种回归方法该参数也需赋值（0或1均可）
    :param float lasso_alpha: lasso回归时的alpha参数，若不传，则使用LassoCV寻找最适合的alpha
    :param float ridge_alpha: ridge回归时的alpha参数，若不传，则使用RidgeCV寻找最适合的alpha
    :return:
    """
    analysis_type_dict = {'asset': list(ASSET_INDEX_DICT.values()),
                          'sw28': SW_CODE_LIST,
                          'zz10': ZZ_CODE_LIST,
                          'barra': BARRA_STYLE_FACTOR}
    reg_dict = {1: 'minimize', 2: 'lasso', 3: 'ridge'}
    index_list = analysis_type_dict[analysis_type]
    reg_metd = reg_dict[reg_method]
    start_date = portfolio_ret.trade_date.min()
    end_date = portfolio_ret.trade_date.max()
    all_factor_list = BOND_7_FACTOR + FAMA_5_FACTOR + BARRA_FULL_FACTOR
    all_factor_list.append('Default')
    factor_list = []
    for factor in index_list:
        if factor in all_factor_list:
            factor_list.append(factor)
    i_df_1 = index_price_related.get_daily_return(index_list, start_date, end_date, order_by='`day`',
                                                  return_only_price_col=False, replace_price=True, replace_tj=True)
    i_df_1['trade_date'] = i_df_1['trade_date'].astype(str)
    i_df_1 = i_df_1.set_index(['trade_date', 'index_code'])['daily_return'].unstack().reset_index()
    i_df_2 = factor_return_related.get_combined_factor_return(factor_list, start_date, end_date)
    i_df_2 = i_df_2.rename(columns={'end_date': 'trade_date'})
    if i_df_1.shape[0] == 0:
        i_df = i_df_2
    elif i_df_2.shape[1] == 1:
        i_df = i_df_1
    else:
        i_df = i_df_1.merge(i_df_2, how='inner', on='trade_date')
    i_df = i_df.dropna()
    df_to_fit = portfolio_ret.merge(i_df, how='inner', on='trade_date').sort_values(by='trade_date')
    result = get_regression_res_by_nav(df_to_fit=df_to_fit,
                                       reg_method=reg_metd,
                                       need_bounds=need_bounds,
                                       need_constraints=need_constraints,
                                       lasso_alpha=lasso_alpha,
                                       ridge_alpha=ridge_alpha)
    result = result[['index_code', 'index_weight']]
    return result